{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# `AnnotateSpans`\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "{func}`tvm.relay.transform.AnnotateSpans` 的作用是为程序添加跨度信息。具体来说，它首先生成程序的文本表示形式，然后将其解析回带有跨度信息的 Relay 抽象语法树（AST）。对模块进行美化打印，然后再将其解析回来，以便为所有 Relay 子表达式建立 spans（范围）和 sources（来源）。这有助于改善程序化构建的模块在下游的错误和调试诊断。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fn (%input: Tensor[(1, 3, 224, 224), float32] /* ty=Tensor[(1, 3, 224, 224), float32] span=GeneratedSource:116:18 */, %aten___convolution_0_weight: Tensor[(64, 3, 7, 7), float32] /* ty=Tensor[(64, 3, 7, 7), float32] span=GeneratedSource:116:28 */, %aten__batch_norm_0_weight: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:117:26 */, %aten__batch_norm_0_bias: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:117:54 */, %aten__batch_norm_0_mean: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:117:80 */, %aten__batch_norm_0_var: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:117:106 */, %aten___convolution_1_weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] span=GeneratedSource:121:22 */, %aten__batch_norm_1_weight: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:122:26 */, %aten__batch_norm_1_bias: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:122:54 */, %aten__batch_norm_1_mean: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:122:80 */, %aten__batch_norm_1_var: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:122:106 */, %aten___convolution_2_weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] span=GeneratedSource:125:22 */, %aten__batch_norm_2_weight: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:126:27 */, %aten__batch_norm_2_bias: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:126:55 */, %aten__batch_norm_2_mean: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:126:81 */, %aten__batch_norm_2_var: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:126:107 */, %aten___convolution_3_weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] span=GeneratedSource:130:24 */, %aten__batch_norm_3_weight: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:131:28 */, %aten__batch_norm_3_bias: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:131:56 */, %aten__batch_norm_3_mean: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:131:82 */, %aten__batch_norm_3_var: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:131:108 */, %aten___convolution_4_weight: Tensor[(64, 64, 3, 3), float32] /* ty=Tensor[(64, 64, 3, 3), float32] span=GeneratedSource:134:24 */, %aten__batch_norm_4_weight: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:135:28 */, %aten__batch_norm_4_bias: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:135:56 */, %aten__batch_norm_4_mean: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:135:82 */, %aten__batch_norm_4_var: Tensor[(64), float32] /* ty=Tensor[(64), float32] span=GeneratedSource:135:108 */, %aten___convolution_5_weight: Tensor[(128, 64, 3, 3), float32] /* ty=Tensor[(128, 64, 3, 3), float32] span=GeneratedSource:139:24 */, %aten__batch_norm_5_weight: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:140:28 */, %aten__batch_norm_5_bias: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:140:56 */, %aten__batch_norm_5_mean: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:140:82 */, %aten__batch_norm_5_var: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:140:108 */, %aten___convolution_6_weight: Tensor[(128, 128, 3, 3), float32] /* ty=Tensor[(128, 128, 3, 3), float32] span=GeneratedSource:143:24 */, %aten__batch_norm_6_weight: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:144:28 */, %aten__batch_norm_6_bias: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:144:56 */, %aten__batch_norm_6_mean: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:144:82 */, %aten__batch_norm_6_var: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:144:108 */, %aten___convolution_7_weight: Tensor[(128, 64, 1, 1), float32] /* ty=Tensor[(128, 64, 1, 1), float32] span=GeneratedSource:145:24 */, %aten__batch_norm_7_weight: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:146:28 */, %aten__batch_norm_7_bias: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:146:56 */, %aten__batch_norm_7_mean: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:146:82 */, %aten__batch_norm_7_var: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:146:108 */, %aten___convolution_8_weight: Tensor[(128, 128, 3, 3), float32] /* ty=Tensor[(128, 128, 3, 3), float32] span=GeneratedSource:151:24 */, %aten__batch_norm_8_weight: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:152:28 */, %aten__batch_norm_8_bias: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:152:56 */, %aten__batch_norm_8_mean: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:152:82 */, %aten__batch_norm_8_var: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:152:108 */, %aten___convolution_9_weight: Tensor[(128, 128, 3, 3), float32] /* ty=Tensor[(128, 128, 3, 3), float32] span=GeneratedSource:155:24 */, %aten__batch_norm_9_weight: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:156:28 */, %aten__batch_norm_9_bias: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:156:56 */, %aten__batch_norm_9_mean: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:156:82 */, %aten__batch_norm_9_var: Tensor[(128), float32] /* ty=Tensor[(128), float32] span=GeneratedSource:156:108 */, %aten___convolution_10_weight: Tensor[(256, 128, 3, 3), float32] /* ty=Tensor[(256, 128, 3, 3), float32] span=GeneratedSource:160:24 */, %aten__batch_norm_10_weight: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:161:28 */, %aten__batch_norm_10_bias: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:161:57 */, %aten__batch_norm_10_mean: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:161:84 */, %aten__batch_norm_10_var: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:161:111 */, %aten___convolution_11_weight: Tensor[(256, 256, 3, 3), float32] /* ty=Tensor[(256, 256, 3, 3), float32] span=GeneratedSource:164:24 */, %aten__batch_norm_11_weight: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:165:28 */, %aten__batch_norm_11_bias: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:165:57 */, %aten__batch_norm_11_mean: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:165:84 */, %aten__batch_norm_11_var: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:165:111 */, %aten___convolution_12_weight: Tensor[(256, 128, 1, 1), float32] /* ty=Tensor[(256, 128, 1, 1), float32] span=GeneratedSource:166:24 */, %aten__batch_norm_12_weight: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:167:28 */, %aten__batch_norm_12_bias: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:167:57 */, %aten__batch_norm_12_mean: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:167:84 */, %aten__batch_norm_12_var: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:167:111 */, %aten___convolution_13_weight: Tensor[(256, 256, 3, 3), float32] /* ty=Tensor[(256, 256, 3, 3), float32] span=GeneratedSource:172:24 */, %aten__batch_norm_13_weight: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:173:28 */, %aten__batch_norm_13_bias: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:173:57 */, %aten__batch_norm_13_mean: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:173:84 */, %aten__batch_norm_13_var: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:173:111 */, %aten___convolution_14_weight: Tensor[(256, 256, 3, 3), float32] /* ty=Tensor[(256, 256, 3, 3), float32] span=GeneratedSource:176:24 */, %aten__batch_norm_14_weight: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:177:28 */, %aten__batch_norm_14_bias: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:177:57 */, %aten__batch_norm_14_mean: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:177:84 */, %aten__batch_norm_14_var: Tensor[(256), float32] /* ty=Tensor[(256), float32] span=GeneratedSource:177:111 */, %aten___convolution_15_weight: Tensor[(512, 256, 3, 3), float32] /* ty=Tensor[(512, 256, 3, 3), float32] span=GeneratedSource:181:24 */, %aten__batch_norm_15_weight: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:182:28 */, %aten__batch_norm_15_bias: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:182:57 */, %aten__batch_norm_15_mean: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:182:84 */, %aten__batch_norm_15_var: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:182:111 */, %aten___convolution_16_weight: Tensor[(512, 512, 3, 3), float32] /* ty=Tensor[(512, 512, 3, 3), float32] span=GeneratedSource:185:24 */, %aten__batch_norm_16_weight: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:186:28 */, %aten__batch_norm_16_bias: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:186:57 */, %aten__batch_norm_16_mean: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:186:84 */, %aten__batch_norm_16_var: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:186:111 */, %aten___convolution_17_weight: Tensor[(512, 256, 1, 1), float32] /* ty=Tensor[(512, 256, 1, 1), float32] span=GeneratedSource:187:24 */, %aten__batch_norm_17_weight: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:188:28 */, %aten__batch_norm_17_bias: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:188:57 */, %aten__batch_norm_17_mean: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:188:84 */, %aten__batch_norm_17_var: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:188:111 */, %aten___convolution_18_weight: Tensor[(512, 512, 3, 3), float32] /* ty=Tensor[(512, 512, 3, 3), float32] span=GeneratedSource:193:24 */, %aten__batch_norm_18_weight: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:194:28 */, %aten__batch_norm_18_bias: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:194:57 */, %aten__batch_norm_18_mean: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:194:84 */, %aten__batch_norm_18_var: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:194:111 */, %aten___convolution_19_weight: Tensor[(512, 512, 3, 3), float32] /* ty=Tensor[(512, 512, 3, 3), float32] span=GeneratedSource:197:24 */, %aten__batch_norm_19_weight: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:198:28 */, %aten__batch_norm_19_bias: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:198:57 */, %aten__batch_norm_19_mean: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:198:84 */, %aten__batch_norm_19_var: Tensor[(512), float32] /* ty=Tensor[(512), float32] span=GeneratedSource:198:111 */, %aten__linear_0_weight: Tensor[(1000, 512), float32] /* ty=Tensor[(1000, 512), float32] span=GeneratedSource:205:23 */, %aten__linear_0_bias: Tensor[(1000), float32] /* ty=Tensor[(1000), float32] span=GeneratedSource:206:20 */) -> Tensor[(1, 1000), float32] {\n",
      "  %0 = nn.conv2d(%input, %aten___convolution_0_weight, strides=[2, 2], padding=[3, 3, 3, 3], channels=64, kernel_size=[7, 7]) /* ty=Tensor[(1, 64, 112, 112), float32] span=GeneratedSource:117:22 */;\n",
      "  %1 = nn.batch_norm(%0, %aten__batch_norm_0_weight, %aten__batch_norm_0_bias, %aten__batch_norm_0_mean, %aten__batch_norm_0_var) /* ty=(Tensor[(1, 64, 112, 112), float32], Tensor[(64), float32], Tensor[(64), float32]) span=GeneratedSource:118:8 */;\n",
      "  %2 = %1.0 /* ty=Tensor[(1, 64, 112, 112), float32] span=GeneratedSource:119:16 */;\n",
      "  %3 = nn.relu(%2) /* ty=Tensor[(1, 64, 112, 112), float32] span=GeneratedSource:120:23 */;\n",
      "  %4 = nn.max_pool2d(%3, pool_size=[3, 3], strides=[2, 2], padding=[1, 1, 1, 1]) /* ty=Tensor[(1, 64, 56, 56), float32] span=GeneratedSource:128:18 */;\n",
      "  %5 = nn.conv2d(%4, %aten___convolution_1_weight, padding=[1, 1, 1, 1], channels=64, kernel_size=[3, 3]) /* ty=Tensor[(1, 64, 56, 56), float32] span=GeneratedSource:122:22 */;\n",
      "  %6 = nn.batch_norm(%5, %aten__batch_norm_1_weight, %aten__batch_norm_1_bias, %aten__batch_norm_1_mean, %aten__batch_norm_1_var) /* ty=(Tensor[(1, 64, 56, 56), float32], Tensor[(64), float32], Tensor[(64), float32]) span=GeneratedSource:123:8 */;\n",
      "  %7 = %6.0 /* ty=Tensor[(1, 64, 56, 56), float32] span=GeneratedSource:124:16 */;\n",
      "  %8 = nn.relu(%7) /* ty=Tensor[(1, 64, 56, 56), float32] span=GeneratedSource:125:18 */;\n",
      "  %9 = nn.conv2d(%8, %aten___convolution_2_weight, padding=[1, 1, 1, 1], channels=64, kernel_size=[3, 3]) /* ty=Tensor[(1, 64, 56, 56), float32] span=GeneratedSource:126:23 */;\n",
      "  %10 = nn.batch_norm(%9, %aten__batch_norm_2_weight, %aten__batch_norm_2_bias, %aten__batch_norm_2_mean, %aten__batch_norm_2_var) /* ty=(Tensor[(1, 64, 56, 56), float32], Tensor[(64), float32], Tensor[(64), float32]) span=GeneratedSource:127:9 */;\n",
      "  %11 = %10.0 /* ty=Tensor[(1, 64, 56, 56), float32] span=GeneratedSource:128:13 */;\n",
      "  %12 = add(%11, %4) /* ty=Tensor[(1, 64, 56, 56), float32] span=GeneratedSource:129:17 */;\n",
      "  %13 = nn.relu(%12) /* ty=Tensor[(1, 64, 56, 56), float32] span=GeneratedSource:137:18 */;\n",
      "  %14 = nn.conv2d(%13, %aten___convolution_3_weight, padding=[1, 1, 1, 1], channels=64, kernel_size=[3, 3]) /* ty=Tensor[(1, 64, 56, 56), float32] span=GeneratedSource:131:23 */;\n",
      "  %15 = nn.batch_norm(%14, %aten__batch_norm_3_weight, %aten__batch_norm_3_bias, %aten__batch_norm_3_mean, %aten__batch_norm_3_var) /* ty=(Tensor[(1, 64, 56, 56), float32], Tensor[(64), float32], Tensor[(64), float32]) span=GeneratedSource:132:9 */;\n",
      "  %16 = %15.0 /* ty=Tensor[(1, 64, 56, 56), float32] span=GeneratedSource:133:17 */;\n",
      "  %17 = nn.relu(%16) /* ty=Tensor[(1, 64, 56, 56), float32] span=GeneratedSource:134:19 */;\n",
      "  %18 = nn.conv2d(%17, %aten___convolution_4_weight, padding=[1, 1, 1, 1], channels=64, kernel_size=[3, 3]) /* ty=Tensor[(1, 64, 56, 56), float32] span=GeneratedSource:135:23 */;\n",
      "  %19 = nn.batch_norm(%18, %aten__batch_norm_4_weight, %aten__batch_norm_4_bias, %aten__batch_norm_4_mean, %aten__batch_norm_4_var) /* ty=(Tensor[(1, 64, 56, 56), float32], Tensor[(64), float32], Tensor[(64), float32]) span=GeneratedSource:136:9 */;\n",
      "  %20 = %19.0 /* ty=Tensor[(1, 64, 56, 56), float32] span=GeneratedSource:137:13 */;\n",
      "  %21 = add(%20, %13) /* ty=Tensor[(1, 64, 56, 56), float32] span=GeneratedSource:138:17 */;\n",
      "  %22 = nn.relu(%21) /* ty=Tensor[(1, 64, 56, 56), float32] span=GeneratedSource:145:19 */;\n",
      "  %23 = nn.conv2d(%22, %aten___convolution_5_weight, strides=[2, 2], padding=[1, 1, 1, 1], channels=128, kernel_size=[3, 3]) /* ty=Tensor[(1, 128, 28, 28), float32] span=GeneratedSource:140:23 */;\n",
      "  %24 = nn.batch_norm(%23, %aten__batch_norm_5_weight, %aten__batch_norm_5_bias, %aten__batch_norm_5_mean, %aten__batch_norm_5_var) /* ty=(Tensor[(1, 128, 28, 28), float32], Tensor[(128), float32], Tensor[(128), float32]) span=GeneratedSource:141:9 */;\n",
      "  %25 = %24.0 /* ty=Tensor[(1, 128, 28, 28), float32] span=GeneratedSource:142:17 */;\n",
      "  %26 = nn.relu(%25) /* ty=Tensor[(1, 128, 28, 28), float32] span=GeneratedSource:143:19 */;\n",
      "  %27 = nn.conv2d(%26, %aten___convolution_6_weight, padding=[1, 1, 1, 1], channels=128, kernel_size=[3, 3]) /* ty=Tensor[(1, 128, 28, 28), float32] span=GeneratedSource:144:23 */;\n",
      "  %28 = nn.batch_norm(%27, %aten__batch_norm_6_weight, %aten__batch_norm_6_bias, %aten__batch_norm_6_mean, %aten__batch_norm_6_var) /* ty=(Tensor[(1, 128, 28, 28), float32], Tensor[(128), float32], Tensor[(128), float32]) span=GeneratedSource:147:9 */;\n",
      "  %29 = nn.conv2d(%22, %aten___convolution_7_weight, strides=[2, 2], padding=[0, 0, 0, 0], channels=128, kernel_size=[1, 1]) /* ty=Tensor[(1, 128, 28, 28), float32] span=GeneratedSource:146:23 */;\n",
      "  %30 = nn.batch_norm(%29, %aten__batch_norm_7_weight, %aten__batch_norm_7_bias, %aten__batch_norm_7_mean, %aten__batch_norm_7_var) /* ty=(Tensor[(1, 128, 28, 28), float32], Tensor[(128), float32], Tensor[(128), float32]) span=GeneratedSource:148:9 */;\n",
      "  %31 = %28.0 /* ty=Tensor[(1, 128, 28, 28), float32] span=GeneratedSource:149:13 */;\n",
      "  %32 = %30.0 /* ty=Tensor[(1, 128, 28, 28), float32] span=GeneratedSource:149:18 */;\n",
      "  %33 = add(%31, %32) /* ty=Tensor[(1, 128, 28, 28), float32] span=GeneratedSource:150:17 */;\n",
      "  %34 = nn.relu(%33) /* ty=Tensor[(1, 128, 28, 28), float32] span=GeneratedSource:158:18 */;\n",
      "  %35 = nn.conv2d(%34, %aten___convolution_8_weight, padding=[1, 1, 1, 1], channels=128, kernel_size=[3, 3]) /* ty=Tensor[(1, 128, 28, 28), float32] span=GeneratedSource:152:23 */;\n",
      "  %36 = nn.batch_norm(%35, %aten__batch_norm_8_weight, %aten__batch_norm_8_bias, %aten__batch_norm_8_mean, %aten__batch_norm_8_var) /* ty=(Tensor[(1, 128, 28, 28), float32], Tensor[(128), float32], Tensor[(128), float32]) span=GeneratedSource:153:9 */;\n",
      "  %37 = %36.0 /* ty=Tensor[(1, 128, 28, 28), float32] span=GeneratedSource:154:17 */;\n",
      "  %38 = nn.relu(%37) /* ty=Tensor[(1, 128, 28, 28), float32] span=GeneratedSource:155:19 */;\n",
      "  %39 = nn.conv2d(%38, %aten___convolution_9_weight, padding=[1, 1, 1, 1], channels=128, kernel_size=[3, 3]) /* ty=Tensor[(1, 128, 28, 28), float32] span=GeneratedSource:156:23 */;\n",
      "  %40 = nn.batch_norm(%39, %aten__batch_norm_9_weight, %aten__batch_norm_9_bias, %aten__batch_norm_9_mean, %aten__batch_norm_9_var) /* ty=(Tensor[(1, 128, 28, 28), float32], Tensor[(128), float32], Tensor[(128), float32]) span=GeneratedSource:157:9 */;\n",
      "  %41 = %40.0 /* ty=Tensor[(1, 128, 28, 28), float32] span=GeneratedSource:158:13 */;\n",
      "  %42 = add(%41, %34) /* ty=Tensor[(1, 128, 28, 28), float32] span=GeneratedSource:159:17 */;\n",
      "  %43 = nn.relu(%42) /* ty=Tensor[(1, 128, 28, 28), float32] span=GeneratedSource:166:19 */;\n",
      "  %44 = nn.conv2d(%43, %aten___convolution_10_weight, strides=[2, 2], padding=[1, 1, 1, 1], channels=256, kernel_size=[3, 3]) /* ty=Tensor[(1, 256, 14, 14), float32] span=GeneratedSource:161:23 */;\n",
      "  %45 = nn.batch_norm(%44, %aten__batch_norm_10_weight, %aten__batch_norm_10_bias, %aten__batch_norm_10_mean, %aten__batch_norm_10_var) /* ty=(Tensor[(1, 256, 14, 14), float32], Tensor[(256), float32], Tensor[(256), float32]) span=GeneratedSource:162:9 */;\n",
      "  %46 = %45.0 /* ty=Tensor[(1, 256, 14, 14), float32] span=GeneratedSource:163:17 */;\n",
      "  %47 = nn.relu(%46) /* ty=Tensor[(1, 256, 14, 14), float32] span=GeneratedSource:164:19 */;\n",
      "  %48 = nn.conv2d(%47, %aten___convolution_11_weight, padding=[1, 1, 1, 1], channels=256, kernel_size=[3, 3]) /* ty=Tensor[(1, 256, 14, 14), float32] span=GeneratedSource:165:23 */;\n",
      "  %49 = nn.batch_norm(%48, %aten__batch_norm_11_weight, %aten__batch_norm_11_bias, %aten__batch_norm_11_mean, %aten__batch_norm_11_var) /* ty=(Tensor[(1, 256, 14, 14), float32], Tensor[(256), float32], Tensor[(256), float32]) span=GeneratedSource:168:9 */;\n",
      "  %50 = nn.conv2d(%43, %aten___convolution_12_weight, strides=[2, 2], padding=[0, 0, 0, 0], channels=256, kernel_size=[1, 1]) /* ty=Tensor[(1, 256, 14, 14), float32] span=GeneratedSource:167:23 */;\n",
      "  %51 = nn.batch_norm(%50, %aten__batch_norm_12_weight, %aten__batch_norm_12_bias, %aten__batch_norm_12_mean, %aten__batch_norm_12_var) /* ty=(Tensor[(1, 256, 14, 14), float32], Tensor[(256), float32], Tensor[(256), float32]) span=GeneratedSource:169:9 */;\n",
      "  %52 = %49.0 /* ty=Tensor[(1, 256, 14, 14), float32] span=GeneratedSource:170:13 */;\n",
      "  %53 = %51.0 /* ty=Tensor[(1, 256, 14, 14), float32] span=GeneratedSource:170:18 */;\n",
      "  %54 = add(%52, %53) /* ty=Tensor[(1, 256, 14, 14), float32] span=GeneratedSource:171:17 */;\n",
      "  %55 = nn.relu(%54) /* ty=Tensor[(1, 256, 14, 14), float32] span=GeneratedSource:179:18 */;\n",
      "  %56 = nn.conv2d(%55, %aten___convolution_13_weight, padding=[1, 1, 1, 1], channels=256, kernel_size=[3, 3]) /* ty=Tensor[(1, 256, 14, 14), float32] span=GeneratedSource:173:23 */;\n",
      "  %57 = nn.batch_norm(%56, %aten__batch_norm_13_weight, %aten__batch_norm_13_bias, %aten__batch_norm_13_mean, %aten__batch_norm_13_var) /* ty=(Tensor[(1, 256, 14, 14), float32], Tensor[(256), float32], Tensor[(256), float32]) span=GeneratedSource:174:9 */;\n",
      "  %58 = %57.0 /* ty=Tensor[(1, 256, 14, 14), float32] span=GeneratedSource:175:17 */;\n",
      "  %59 = nn.relu(%58) /* ty=Tensor[(1, 256, 14, 14), float32] span=GeneratedSource:176:19 */;\n",
      "  %60 = nn.conv2d(%59, %aten___convolution_14_weight, padding=[1, 1, 1, 1], channels=256, kernel_size=[3, 3]) /* ty=Tensor[(1, 256, 14, 14), float32] span=GeneratedSource:177:23 */;\n",
      "  %61 = nn.batch_norm(%60, %aten__batch_norm_14_weight, %aten__batch_norm_14_bias, %aten__batch_norm_14_mean, %aten__batch_norm_14_var) /* ty=(Tensor[(1, 256, 14, 14), float32], Tensor[(256), float32], Tensor[(256), float32]) span=GeneratedSource:178:9 */;\n",
      "  %62 = %61.0 /* ty=Tensor[(1, 256, 14, 14), float32] span=GeneratedSource:179:13 */;\n",
      "  %63 = add(%62, %55) /* ty=Tensor[(1, 256, 14, 14), float32] span=GeneratedSource:180:17 */;\n",
      "  %64 = nn.relu(%63) /* ty=Tensor[(1, 256, 14, 14), float32] span=GeneratedSource:187:19 */;\n",
      "  %65 = nn.conv2d(%64, %aten___convolution_15_weight, strides=[2, 2], padding=[1, 1, 1, 1], channels=512, kernel_size=[3, 3]) /* ty=Tensor[(1, 512, 7, 7), float32] span=GeneratedSource:182:23 */;\n",
      "  %66 = nn.batch_norm(%65, %aten__batch_norm_15_weight, %aten__batch_norm_15_bias, %aten__batch_norm_15_mean, %aten__batch_norm_15_var) /* ty=(Tensor[(1, 512, 7, 7), float32], Tensor[(512), float32], Tensor[(512), float32]) span=GeneratedSource:183:9 */;\n",
      "  %67 = %66.0 /* ty=Tensor[(1, 512, 7, 7), float32] span=GeneratedSource:184:17 */;\n",
      "  %68 = nn.relu(%67) /* ty=Tensor[(1, 512, 7, 7), float32] span=GeneratedSource:185:19 */;\n",
      "  %69 = nn.conv2d(%68, %aten___convolution_16_weight, padding=[1, 1, 1, 1], channels=512, kernel_size=[3, 3]) /* ty=Tensor[(1, 512, 7, 7), float32] span=GeneratedSource:186:23 */;\n",
      "  %70 = nn.batch_norm(%69, %aten__batch_norm_16_weight, %aten__batch_norm_16_bias, %aten__batch_norm_16_mean, %aten__batch_norm_16_var) /* ty=(Tensor[(1, 512, 7, 7), float32], Tensor[(512), float32], Tensor[(512), float32]) span=GeneratedSource:189:9 */;\n",
      "  %71 = nn.conv2d(%64, %aten___convolution_17_weight, strides=[2, 2], padding=[0, 0, 0, 0], channels=512, kernel_size=[1, 1]) /* ty=Tensor[(1, 512, 7, 7), float32] span=GeneratedSource:188:23 */;\n",
      "  %72 = nn.batch_norm(%71, %aten__batch_norm_17_weight, %aten__batch_norm_17_bias, %aten__batch_norm_17_mean, %aten__batch_norm_17_var) /* ty=(Tensor[(1, 512, 7, 7), float32], Tensor[(512), float32], Tensor[(512), float32]) span=GeneratedSource:190:9 */;\n",
      "  %73 = %70.0 /* ty=Tensor[(1, 512, 7, 7), float32] span=GeneratedSource:191:13 */;\n",
      "  %74 = %72.0 /* ty=Tensor[(1, 512, 7, 7), float32] span=GeneratedSource:191:18 */;\n",
      "  %75 = add(%73, %74) /* ty=Tensor[(1, 512, 7, 7), float32] span=GeneratedSource:192:17 */;\n",
      "  %76 = nn.relu(%75) /* ty=Tensor[(1, 512, 7, 7), float32] span=GeneratedSource:200:18 */;\n",
      "  %77 = nn.conv2d(%76, %aten___convolution_18_weight, padding=[1, 1, 1, 1], channels=512, kernel_size=[3, 3]) /* ty=Tensor[(1, 512, 7, 7), float32] span=GeneratedSource:194:23 */;\n",
      "  %78 = nn.batch_norm(%77, %aten__batch_norm_18_weight, %aten__batch_norm_18_bias, %aten__batch_norm_18_mean, %aten__batch_norm_18_var) /* ty=(Tensor[(1, 512, 7, 7), float32], Tensor[(512), float32], Tensor[(512), float32]) span=GeneratedSource:195:9 */;\n",
      "  %79 = %78.0 /* ty=Tensor[(1, 512, 7, 7), float32] span=GeneratedSource:196:17 */;\n",
      "  %80 = nn.relu(%79) /* ty=Tensor[(1, 512, 7, 7), float32] span=GeneratedSource:197:19 */;\n",
      "  %81 = nn.conv2d(%80, %aten___convolution_19_weight, padding=[1, 1, 1, 1], channels=512, kernel_size=[3, 3]) /* ty=Tensor[(1, 512, 7, 7), float32] span=GeneratedSource:198:23 */;\n",
      "  %82 = nn.batch_norm(%81, %aten__batch_norm_19_weight, %aten__batch_norm_19_bias, %aten__batch_norm_19_mean, %aten__batch_norm_19_var) /* ty=(Tensor[(1, 512, 7, 7), float32], Tensor[(512), float32], Tensor[(512), float32]) span=GeneratedSource:199:9 */;\n",
      "  %83 = %82.0 /* ty=Tensor[(1, 512, 7, 7), float32] span=GeneratedSource:200:13 */;\n",
      "  %84 = add(%83, %76) /* ty=Tensor[(1, 512, 7, 7), float32] span=GeneratedSource:201:17 */;\n",
      "  %85 = nn.relu(%84) /* ty=Tensor[(1, 512, 7, 7), float32] span=GeneratedSource:202:32 */;\n",
      "  %86 = nn.adaptive_avg_pool2d(%85, output_size=[1, 1]) /* ty=Tensor[(1, 512, 1, 1), float32] span=GeneratedSource:203:17 */;\n",
      "  %87 = reshape(%86, newshape=[0, -1, 1, 1]) /* ty=Tensor[(1, 512, 1, 1), float32] span=GeneratedSource:204:17 */;\n",
      "  %88 = squeeze(%87, axis=[2, 3]) /* ty=Tensor[(1, 512), float32] span=GeneratedSource:205:18 */;\n",
      "  %89 = nn.dense(%88, %aten__linear_0_weight, units=None) /* ty=Tensor[(1, 1000), float32] span=GeneratedSource:206:15 */;\n",
      "  nn.bias_add(%89, %aten__linear_0_bias, axis=-1) /* ty=Tensor[(1, 1000), float32] span=GeneratedSource:116:3 */\n",
      "} /* ty=fn (Tensor[(1, 3, 224, 224), float32], Tensor[(64, 3, 7, 7), float32], Tensor[(64), float32], Tensor[(64), float32], Tensor[(64), float32], Tensor[(64), float32], Tensor[(64, 64, 3, 3), float32], Tensor[(64), float32], Tensor[(64), float32], Tensor[(64), float32], Tensor[(64), float32], Tensor[(64, 64, 3, 3), float32], Tensor[(64), float32], Tensor[(64), float32], Tensor[(64), float32], Tensor[(64), float32], Tensor[(64, 64, 3, 3), float32], Tensor[(64), float32], Tensor[(64), float32], Tensor[(64), float32], Tensor[(64), float32], Tensor[(64, 64, 3, 3), float32], Tensor[(64), float32], Tensor[(64), float32], Tensor[(64), float32], Tensor[(64), float32], Tensor[(128, 64, 3, 3), float32], Tensor[(128), float32], Tensor[(128), float32], Tensor[(128), float32], Tensor[(128), float32], Tensor[(128, 128, 3, 3), float32], Tensor[(128), float32], Tensor[(128), float32], Tensor[(128), float32], Tensor[(128), float32], Tensor[(128, 64, 1, 1), float32], Tensor[(128), float32], Tensor[(128), float32], Tensor[(128), float32], Tensor[(128), float32], Tensor[(128, 128, 3, 3), float32], Tensor[(128), float32], Tensor[(128), float32], Tensor[(128), float32], Tensor[(128), float32], Tensor[(128, 128, 3, 3), float32], Tensor[(128), float32], Tensor[(128), float32], Tensor[(128), float32], Tensor[(128), float32], Tensor[(256, 128, 3, 3), float32], Tensor[(256), float32], Tensor[(256), float32], Tensor[(256), float32], Tensor[(256), float32], Tensor[(256, 256, 3, 3), float32], Tensor[(256), float32], Tensor[(256), float32], Tensor[(256), float32], Tensor[(256), float32], Tensor[(256, 128, 1, 1), float32], Tensor[(256), float32], Tensor[(256), float32], Tensor[(256), float32], Tensor[(256), float32], Tensor[(256, 256, 3, 3), float32], Tensor[(256), float32], Tensor[(256), float32], Tensor[(256), float32], Tensor[(256), float32], Tensor[(256, 256, 3, 3), float32], Tensor[(256), float32], Tensor[(256), float32], Tensor[(256), float32], Tensor[(256), float32], Tensor[(512, 256, 3, 3), float32], Tensor[(512), float32], Tensor[(512), float32], Tensor[(512), float32], Tensor[(512), float32], Tensor[(512, 512, 3, 3), float32], Tensor[(512), float32], Tensor[(512), float32], Tensor[(512), float32], Tensor[(512), float32], Tensor[(512, 256, 1, 1), float32], Tensor[(512), float32], Tensor[(512), float32], Tensor[(512), float32], Tensor[(512), float32], Tensor[(512, 512, 3, 3), float32], Tensor[(512), float32], Tensor[(512), float32], Tensor[(512), float32], Tensor[(512), float32], Tensor[(512, 512, 3, 3), float32], Tensor[(512), float32], Tensor[(512), float32], Tensor[(512), float32], Tensor[(512), float32], Tensor[(1000, 512), float32], Tensor[(1000), float32]) -> Tensor[(1, 1000), float32] span=GeneratedSource:115:1 */\n"
     ]
    }
   ],
   "source": [
    "import torchvision\n",
    "import torch\n",
    "from tvm import relay\n",
    "model = torchvision.models.resnet18().eval()\n",
    "inp = torch.randn([1, 3, 224, 224])\n",
    "trace = torch.jit.trace(model, inp).eval()\n",
    "mod, _ = relay.frontend.from_pytorch(\n",
    "    trace, [(\"input\", inp.shape)], use_parser_friendly_name=True\n",
    ")\n",
    "mod = relay.transform.AnnotateSpans()(mod)\n",
    "print(mod[\"main\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>), float32] span<span style=\"color: #AA22FF; font-weight: bold\">=</span>GeneratedSource:<span style=\"color: #008000\">3</span>:<span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>), float32] {\n",
       "  clip(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x, a_min<span style=\"color: #AA22FF; font-weight: bold\">=-</span>inff, a_max<span style=\"color: #AA22FF; font-weight: bold\">=</span>inff) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>), float32] span<span style=\"color: #AA22FF; font-weight: bold\">=</span>GeneratedSource:<span style=\"color: #008000\">3</span>:<span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import tvm\n",
    "x = relay.var(\"x\", shape=(3, 4), dtype=\"float32\")\n",
    "y = relay.clip(x, -np.inf, np.inf)\n",
    "\n",
    "f = relay.Function([x], y)\n",
    "mod = tvm.IRModule.from_expr(f)\n",
    "\n",
    "mod = relay.transform.AnnotateSpans()(mod)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[11:36:37] /media/pc/data/lxw/ai/tvm/src/target/target_kind.cc:171: Warning: Unable to detect CUDA version, default to \"-arch=sm_50\" instead\n"
     ]
    }
   ],
   "source": [
    "import tvm\n",
    "import tvm.relay as relay\n",
    "from tvm.relay import testing\n",
    "import tvm.testing\n",
    "\n",
    "\n",
    "def test_annotate_spans_compatibility():\n",
    "    data = relay.var(\"data\", relay.TensorType((1, 3, 64, 64), \"float32\"))\n",
    "    weight = relay.var(\"weight\")\n",
    "\n",
    "    bn_gamma = relay.var(\"bn_gamma\")\n",
    "    bn_beta = relay.var(\"bn_beta\")\n",
    "    bn_mmean = relay.var(\"bn_mean\")\n",
    "    bn_mvar = relay.var(\"bn_var\")\n",
    "\n",
    "    simple_net = relay.nn.conv2d(\n",
    "        data=data, weight=weight, kernel_size=(3, 3), channels=3, padding=(1, 1)\n",
    "    )\n",
    "    simple_net = relay.nn.batch_norm(simple_net, bn_gamma, bn_beta, bn_mmean, bn_mvar)[0]\n",
    "    simple_net = relay.Function(relay.analysis.free_vars(simple_net), simple_net)\n",
    "\n",
    "    module, params = testing.create_workload(simple_net)\n",
    "\n",
    "    # Apply some simple passes to legalize the IR.\n",
    "    with tvm.transform.PassContext(opt_level=0):\n",
    "        module, params = relay.optimize(\n",
    "            module, target=tvm.testing.enabled_targets()[0][0], params=params\n",
    "        )\n",
    "\n",
    "    seq = tvm.transform.Sequential([relay.transform.AnnotateSpans(), relay.transform.DefuseOps()])\n",
    "    with tvm.transform.PassContext(opt_level=3):\n",
    "        module = seq(module)\n",
    "\n",
    "test_annotate_spans_compatibility()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "import tvm.relay\n",
    "from tvm.relay import op\n",
    "from tvm.ir.instrument import PassTimingInstrument, pass_instrument\n",
    "\n",
    "def get_test_model():\n",
    "    x, y, z = [tvm.relay.var(c, shape=(3, 4), dtype=\"float32\") for c in \"xyz\"]\n",
    "    e1 = op.add(x, y)\n",
    "    e2 = op.subtract(x, z)\n",
    "    e3 = op.multiply(e1, e1 / e2)\n",
    "    return tvm.IRModule.from_expr(e3 + e2)\n",
    "\n",
    "def test_pass_timing_instrument():\n",
    "    pass_timing = PassTimingInstrument()\n",
    "\n",
    "    # Override current PassContext's instruments\n",
    "    tvm.transform.PassContext.current().override_instruments([pass_timing])\n",
    "\n",
    "    mod = get_test_model()\n",
    "    mod = tvm.relay.transform.AnnotateSpans()(mod)\n",
    "    mod = tvm.relay.transform.ToANormalForm()(mod)\n",
    "    mod = tvm.relay.transform.InferType()(mod)\n",
    "\n",
    "    profiles = pass_timing.render()\n",
    "    assert \"AnnotateSpans\" in profiles\n",
    "    assert \"ToANormalForm\" in profiles\n",
    "    assert \"InferType\" in profiles\n",
    "\n",
    "    # Reset current PassContext's instruments to None\n",
    "    tvm.transform.PassContext.current().override_instruments(None)\n",
    "\n",
    "    mod = get_test_model()\n",
    "    mod = tvm.relay.transform.AnnotateSpans()(mod)\n",
    "    mod = tvm.relay.transform.ToANormalForm()(mod)\n",
    "    mod = tvm.relay.transform.InferType()(mod)\n",
    "\n",
    "    profiles = pass_timing.render()\n",
    "    assert profiles == \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "xxx",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
